import asyncio
import cv2
import numpy as np
import ffmpeg
import pyaudio
from multiprocessing import Process

ErrorStatus = False


# def videoGet():
#     inputs = ffmpeg.input("rtmp://139.155.181.21/live/01")
#     video = inputs.video
#     videoStream = video.output(
#         "pipe:", format="rawvideo", pix_fmt="bgr24"
#     ).run_async(pipe_stdout=True)
#     while True:
#         dataRaw = videoStream.stdout.read(720 * 1280 * 3)
#         frame = np.frombuffer(dataRaw, np.uint8).reshape((720, 1280, 3))
#         cv2.imshow("test", frame)
#         if cv2.waitKey(1) & 0xff == ord('q'):
#             break
#     cv2.destroyAllWindows()
#     videoStream.terminate()
#
#
# def audioGet():
#     p = pyaudio.PyAudio()
#     stream = p.open(
#         format=pyaudio.paInt16,
#         channels=2,
#         rate=48000,
#         output=True,
#     )
#     inputs = ffmpeg.input("rtmp://139.155.181.21/live/01")
#     audio = inputs.audio
#     audioStream = audio.output(
#         "pipe:", format="wav",
#     ).run_async(pipe_stdout=True)
#     while not ErrorStatus:
#         dataRaw = audioStream.stdout.read(1024)
#         stream.write(dataRaw)
#     stream.close()
#     audioStream.terminate()
#
#
# if __name__ == '__main__':
#     proc1 = Process(target=videoGet(), daemon=True)
#     proc2 = Process(target=audioGet(), daemon=True)
#     proc1.start()
#     proc2.start()
#     proc1.join()
#     proc2.join()

async def __videoGetTask(videoStream):
    global ErrorStatus
    while not ErrorStatus:
        dataRaw = await videoStream.stdout.read(720 * 1280 * 3)
        bytesSize = len(dataRaw)
        if bytesSize == 720 * 1280 * 3:
            frame = np.frombuffer(dataRaw, np.uint8).reshape((720, 1280, 3)
            cv2.imshow("test", frame)
            if cv2.waitKey(1) & 0xff == ord('q'):
                break
    ErrorStatus = True


#
# async def __audioGetTask(audioStream):
#     global ErrorStatus
#     while not ErrorStatus:
#         dataRaw = audioStream.stdout.read(1024)
#         stream.write(dataRaw)
#         await asyncio.sleep(0)
#     ErrorStatus = True
#
#
async def main():

    videoStream = await asyncio.subprocess.create_subprocess_shell(
        "ffmpeg -f flv -i rtmp://139.155.181.21/live/01 -an -f rawvideo -pix_fmt bgr24 -",
        stdout=asyncio.subprocess.PIPE,
        shell=True
    )
    # inputs = ffmpeg.input("rtmp://139.155.181.21/live/01")
    # videoStream = video.output(
    #     "pipe:", format="rawvideo", pix_fmt="bgr24"
    # ).run_async(pipe_stdout=True)
    await asyncio.gather(__videoGetTask(videoStream))

#
asyncio.run(main())


# #
# inputs = ffmpeg.input("rtmp://139.155.181.21/live/01")
# audio = inputs.audio
# video = inputs.video
# audioStream = audio.output(
#     "pipe:", format="wav",
# ).run_async(pipe_stdout=True)
# videoStream = video.output(
#     "pipe:", format="rawvideo", pix_fmt="bgr24"
# ).run_async(pipe_stdout=True)


# def callback(in_data, frame_count, time_info, status):
#     data = audioStream.stdout.read(1024)
#     return data, pyaudio.paContinue


# p = pyaudio.PyAudio()
# stream = p.open(
#     format=pyaudio.paInt16,
#     channels=2,
#     rate=48000,
#     output=True,
# )
#
# while True:
#     in_bytes = audioStream.stdout.read(1024)
#     stream.write(in_bytes)
#     if not in_bytes:
#         break
# frame = np.frombuffer(in_bytes, np.uint8).reshape((720, 1280, 3))
# cv2.imshow("test", frame)
# if cv2.waitKey(1) & 0xff == ord('q'):
#     break

# videoStream.terminate()

# ErrorStatus = False
# #
# cap = cv2.VideoCapture("rtmp://139.155.181.21/live/03")
# while True:
#     try:
#         ret, frame = cap.read()
#         if ret:
#             cv2.imshow("test", cv2.flip(frame, 1))
#             cv2.waitKey(1)
#     except KeyboardInterrupt:
#         break
#
# cap.release()
# async def videoGetTask(saveBuffer: asyncio.Queue):
#     global ErrorStatus
#     videoStream = await asyncio.create_subprocess_shell(
#         "ffmpeg -f flv -i rtmp://139.155.181.21/live/01 -vcodec h264 -f h264 -",
#         stdout=asyncio.subprocess.PIPE,
#         stderr=asyncio.subprocess.DEVNULL,
#         shell=True
#     )
#     while not ErrorStatus:
#         try:
#             rawData = await videoStream.stdout.readexactly(4096)
#             await saveBuffer.put(rawData)
#         except Exception as e:
#             repr(e)
#             break
#     videoStream.terminate()
#
#
# async def videoPutTask(videoStream, saveBuffer: asyncio.Queue):
#     global ErrorStatus
#     while not ErrorStatus:
#         try:
#             if not saveBuffer.empty():
#                 rawInData = await saveBuffer.get()
#                 # print("write", rawInData)
#                 videoStream.stdin.write(rawInData)
#                 await videoStream.stdin.drain()
#             else:
#                 await asyncio.sleep(0)
#         except Exception as e:
#             repr(e)
#             break
#     videoStream.terminate()
#
#
# async def videoShowTask(videoStream):
#     global ErrorStatus
#     while not ErrorStatus:
#         try:
#             rawData = await videoStream.stdout.readexactly(480 * 640 * 3)
#             frame = np.frombuffer(rawData, np.uint8)
#             frame = frame.reshape((480, 640, 3))
#             cv2.imshow('test', frame)
#             cv2.waitKey(1)
#             await asyncio.sleep(0)
#         except Exception as e:
#             repr(e)
#             break
#     videoStream.terminate()
#
#
# async def main():
#     videoStream = await asyncio.create_subprocess_shell(
#         f"ffmpeg -f h264 -i - -vcodec h264 -pix_fmt bgr24 -f rawvideo pipe:",
#         stdout=asyncio.subprocess.PIPE,
#         stdin=asyncio.subprocess.PIPE,
#         stderr=asyncio.subprocess.DEVNULL
#     )
#
#     EncodeBuffer = asyncio.Queue(50)
#
#     getTask = asyncio.create_task(videoGetTask(EncodeBuffer))
#     putTask = asyncio.create_task(videoPutTask(videoStream, EncodeBuffer))
#     showTask = asyncio.create_task(videoShowTask(videoStream))
#
#     await asyncio.gather(getTask, putTask, showTask)
#
# asyncio.run(main())
